experimental/conv/implicit_gemm_cuda.py (5)

134-138: Stale template-parameter comments — BLOCK_N is 64, not 32.

The comments on lines 135 and 138 say BLOCK_N = 32 and WARPS_N = 2, but every actual instantiation (lines 554, 559) uses BLOCK_N=64, WARPS_N=4. Similarly, the comment on line 423 says 64 * 32 * 4 = 8192 bytes when the real footprint is 64 * 64 * 4 = 16384 bytes. The code is correct (it's fully parameterized), but these stale comments will mislead anyone reading the kernel.

---

255-319: Quantized A-tile load: FP4 block size is implicitly coupled to BLOCK_K.

The quantize-dequantize path computes one block_max per warp-row (i.e., over BLOCK_K elements via warp_reduce_max). This means the FP4 quantization block size is always exactly BLOCK_K, which only works correctly because BLOCK_K == fp4_block_size for both supported configs. If a future config changes BLOCK_K independently of fp4_block_size, quantization granularity will silently break. Worth a brief comment or a static_assert in the kernel:

static_assert(BLOCK_K == 128 || BLOCK_K == 256, "BLOCK_K must match fp4_block_size");

---

578-591: verbose=True will spam build logs on every first invocation.

For an experimental module this is fine during development, but consider gating it behind an environment variable or defaulting to False so downstream users don't get unexpected compiler output.

---

641-646: Input validation uses bare assert, which is stripped under python -O.

The assert statements on lines 643 and 646 will be silently removed when Python runs with optimizations enabled. For a CUDA kernel wrapper, invalid shapes reaching the kernel could cause silent corruption or hard crashes. Consider using explicit checks:

Proposed fix

-    assert x.ndim == 5 and w.ndim == 5
+    if x.ndim != 5 or w.ndim != 5:
+        raise ValueError(f"Expected 5D tensors, got x.ndim={x.ndim}, w.ndim={w.ndim}")
     n_batch, cin, d, h, w_in = x.shape
     cout, cin_w, kd, kh, kw = w.shape
-    assert cin_w == cin
+    if cin_w != cin:
+        raise ValueError(f"Input channels mismatch: x has {cin}, w has {cin_w}")

---

663-667: All inputs are cast to .float() (FP32) — potential unnecessary memory doubling.

If inputs are already FP32, the .float().contiguous() calls are cheap. But if inputs arrive as BF16 (common for the use-case described), this silently doubles memory. The docstring says "BF16 WMMA" but the kernel actually consumes FP32 global-memory inputs and converts to BF16 only in shared memory. This is worth a brief comment so users understand the kernel is not end-to-end BF16 in global memory.

experimental/conv/implicit_gemm_binding.cpp (1)

1-20: Duplicate license header block.

The file contains two license header blocks (lines 1-16 and lines 18-19). Consider removing the duplicate.

Proposed fix

-// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-// SPDX-License-Identifier: Apache-2.0
-
 `#include` <torch/extension.h>

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@experimental/conv/implicit_gemm_binding.cpp` around lines 1 - 20, Remove the
duplicated SPDX/license header at the top of the file by keeping a single
canonical license block and deleting the second repeated block (the extra lines
containing "SPDX-FileCopyrightText" and "SPDX-License-Identifier" that appear
immediately after the first header). Ensure only one license header remains at
the top of implicit_gemm_binding.cpp and adjust surrounding whitespace so the
file begins with a single, correctly formatted header.

experimental/conv/implicit_gemm_kernel.cu (1)

1-32: Duplicate license header block.

Similar to the binding file, this contains duplicate license blocks (lines 1-16 and lines 18-31).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@experimental/conv/implicit_gemm_kernel.cu` around lines 1 - 32, The file
contains two repeated license header blocks starting with
"SPDX-FileCopyrightText" and "SPDX-License-Identifier" (the duplicate Apache-2.0
header); remove the redundant block so only a single license header remains at
the top of implicit_gemm_kernel.cu, keeping one copy of the SPDX lines and the
Apache-2.0 boilerplate and deleting the second copy.

modelopt/torch/quantization/nn/modules/quant_conv.py (1)

125-162: Bypasses parent's quantize_weight() context manager pattern.

The parent class QuantLinearConvBase.forward uses a with self.quantize_weight(): context that sets _enable_weight_quantization. This custom forward bypasses that pattern entirely by directly calling _nvfp4_quantize_weight_along_k.

While this may be intentional for the NVFP4 path (since weight quantization is handled differently), consider adding a brief comment explaining why the context manager pattern is not used here, to prevent confusion for future maintainers.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@modelopt/torch/quantization/nn/modules/quant_conv.py` around lines 125 - 162,
The NVFP4 forward bypasses the parent QuantLinearConvBase.forward's with
self.quantize_weight() context and directly calls _nvfp4_quantize_weight_along_k
without setting _enable_weight_quantization; add a short comment in forward
(near the call to _nvfp4_quantize_weight_along_k) that explains this is
intentional because NVFP4 performs weight quantization differently (so the
quantize_weight() context and its _enable_weight_quantization flag must not be
used), and reference both quantize_weight() and _enable_weight_quantization to
help future maintainers understand why the parent pattern is omitted.

experimental/conv/test_implicit_gemm.py (1)

26-31: Consider adding a CUDA availability check in the fixture.

The fixture imports the CUDA module unconditionally. If tests are run on a system without CUDA, the import will fail with a non-obvious error. Consider wrapping with a CUDA availability check.

Proposed enhancement

 `@pytest.fixture`(scope="module")
 def cuda_conv3d():
     """Import and return the CUDA implicit GEMM conv3d function."""
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
     from experimental.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda

     return conv3d_implicit_gemm_cuda

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@experimental/conv/test_implicit_gemm.py` around lines 26 - 31, The fixture
cuda_conv3d currently imports conv3d_implicit_gemm_cuda unconditionally which
will error on machines without CUDA; modify the fixture to first check CUDA
availability (e.g., via pytest.importorskip or torch.cuda.is_available()) and
skip the tests if CUDA is not present, then import or return
conv3d_implicit_gemm_cuda only when CUDA is available so the test suite fails
gracefully on non-CUDA hosts.

experimental/conv/implicit_gemm_kernel.cu (1)

1-32: Duplicate license header.

The license header appears twice (lines 1-16 in block comment style and lines 18-31 in line comment style). Remove one of them.

Proposed fix: Remove the duplicate header

 /*
  * SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
  * http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */

-// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-// SPDX-License-Identifier: Apache-2.0
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
 // Conv3D Implicit GEMM with BF16 WMMA Tensor Cores and optional fused FP4 quantization.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@experimental/conv/implicit_gemm_kernel.cu` around lines 1 - 32, The file
contains a duplicated license header: a block comment header (/* ... */)
followed immediately by the same SPDX/license text in line-comment form (//
...). Remove one of the two headers so only a single canonical license header
remains at the top of implicit_gemm_kernel.cu (either keep the block comment or
the line-comment version), ensuring the SPDX identifiers and full license text
are preserved in the remaining header.

experimental/conv/implicit_gemm_cuda.py (1)

182-212: Add block_size validation to fp4_fake_quant for consistency with conv3d_implicit_gemm_cuda.

The fp4_fake_quant function accepts any block_size value without validation, while conv3d_implicit_gemm_cuda restricts fp4_block_size to {16, 32, 64, 128, 256}. Since the underlying CUDA kernel and documentation confirm only these five sizes are supported, add explicit validation to fail fast on invalid inputs and match the pattern used elsewhere in the module.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@experimental/conv/implicit_gemm_cuda.py` around lines 182 - 212,
fp4_fake_quant currently accepts any block_size but the CUDA kernel only
supports block sizes {16, 32, 64, 128, 256}; add a fast-fail validation at the
start of fp4_fake_quant (similar to conv3d_implicit_gemm_cuda's fp4_block_size
checks) that raises a clear error if block_size is not one of those five values,
referencing the allowed set in the message so callers know valid options.